import { SyntaxKind } from "ts-morph"

import type { Transformer } from "~/utils/transformers"

export const transformTwPrefix: Transformer = async ({ sourceFile, config }) => {
  if (!config.tailwind.prefix) {
    return sourceFile
  }

  // Find the cva function calls.
  sourceFile
    .getDescendantsOfKind(SyntaxKind.CallExpression)
    .filter((node) => node.getExpression().getText() === "cva")
    .forEach((node) => {
      // cva(base, ...)
      if (node.getArguments()[0]?.isKind(SyntaxKind.StringLiteral)) {
        const defaultClassNames = node.getArguments()[0]
        if (defaultClassNames) {
          defaultClassNames.replaceWithText(
            `"${applyPrefix(defaultClassNames.getText()?.replace(/"/g, ""), config.tailwind.prefix)}"`
          )
        }
      }

      // cva(..., { variants: { ... } })
      if (node.getArguments()[1]?.isKind(SyntaxKind.ObjectLiteralExpression)) {
        node
          .getArguments()[1]
          ?.getDescendantsOfKind(SyntaxKind.PropertyAssignment)
          .find((node) => node.getName() === "variants")
          ?.getDescendantsOfKind(SyntaxKind.PropertyAssignment)
          .forEach((node) => {
            node.getDescendantsOfKind(SyntaxKind.PropertyAssignment).forEach((node) => {
              const classNames = node.getInitializerIfKind(SyntaxKind.StringLiteral)
              if (classNames) {
                classNames?.replaceWithText(
                  `"${applyPrefix(classNames.getText()?.replace(/"/g, ""), config.tailwind.prefix)}"`
                )
              }
            })
          })
      }
    })

  // Find all jsx attributes with the name class.
  sourceFile.getDescendantsOfKind(SyntaxKind.JsxAttribute).forEach((node) => {
    if (node.getName() === "class") {
      // class="..."
      if (node.getInitializer()?.isKind(SyntaxKind.StringLiteral)) {
        const value = node.getInitializer()
        if (value) {
          value.replaceWithText(
            `"${applyPrefix(value.getText()?.replace(/"/g, ""), config.tailwind.prefix)}"`
          )
        }
      }

      // class={...}
      if (node.getInitializer()?.isKind(SyntaxKind.JsxExpression)) {
        // Check if it's a call to cn().
        const callExpression = node
          .getInitializer()
          ?.getDescendantsOfKind(SyntaxKind.CallExpression)
          .find((node) => node.getExpression().getText() === "cn")
        if (callExpression) {
          // Loop through the arguments.
          callExpression.getArguments().forEach((node) => {
            if (
              node.isKind(SyntaxKind.ConditionalExpression) ||
              node.isKind(SyntaxKind.BinaryExpression)
            ) {
              node.getChildrenOfKind(SyntaxKind.StringLiteral).forEach((node) => {
                node.replaceWithText(
                  `"${applyPrefix(node.getText()?.replace(/"/g, ""), config.tailwind.prefix)}"`
                )
              })
            }

            if (node.isKind(SyntaxKind.StringLiteral)) {
              node.replaceWithText(
                `"${applyPrefix(node.getText()?.replace(/"/g, ""), config.tailwind.prefix)}"`
              )
            }
          })
        }
      }
    }
  })

  return sourceFile
}

export function applyPrefix(input: string, prefix: string = "") {
  const classNames = input.split(" ")
  const prefixed: string[] = []
  for (const className of classNames) {
    const [variant, value, modifier] = splitClassName(className)
    if (variant) {
      modifier
        ? prefixed.push(`${variant}:${prefix}${value}/${modifier}`)
        : prefixed.push(`${variant}:${prefix}${value}`)
    } else {
      modifier ? prefixed.push(`${prefix}${value}/${modifier}`) : prefixed.push(`${prefix}${value}`)
    }
  }
  return prefixed.join(" ")
}

// Splits a class into variant-name-alpha.
// eg. hover:bg-primary-100 -> [hover, bg-primary, 100]
export function splitClassName(className: string): (string | null)[] {
  if (!className.includes("/") && !className.includes(":")) {
    return [null, className, null]
  }

  const parts: (string | null)[] = []
  // First we split to find the alpha.
  const [rest, alpha] = className.split("/")

  // Check if rest has a colon.
  if (!rest?.includes(":")) {
    return [null, rest!, alpha!]
  }

  // Next we split the rest by the colon.
  const split = rest?.split(":")

  // We take the last item from the split as the name.
  const name = split.pop()

  // We glue back the rest of the split.
  const variant = split.join(":")

  // Finally we push the variant, name and alpha.
  parts.push(variant ?? null, name ?? null, alpha ?? null)

  return parts
}
